#include "multionlineextract.h"

#include <errno.h>
#include <math.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>

#define maybe_fprintf(...)              \
  do                                    \
    {                                   \
      if (! quiet)                      \
        {                               \
          fprintf (__VA_ARGS__);        \
        }                               \
    }                                   \
  while (0)

static float
do_block (MultiModel*   model,
          uint64_t      t,
          float         eta,
          float         rho,
          uint64_t*     last_t,
          float*        dloga,
          float*        dc,
          float*        dd,
          float*        logbeta,
          float*        pz,
          int           clamp,
          const Rating* ratings,
          unsigned int  n_ratings,
          bool          test_only)
{
  return blockopt (t,
                   eta,
                   rho,
                   last_t,
                   model->n_workers,
                   dloga,
                   dc,
                   dd,
                   model->n_items,
                   model->loga,
                   model->gamma,
                   model->rank,
                   logbeta,
                   model->c,
                   model->n_labels,
                   model->d,
                   pz,
                   model->logpriorz,
                   clamp,
                   ratings,
                   n_ratings,
                   model->priorloga,
                   model->priorlogbeta,
                   model->priorc,
                   model->priord,
                   test_only);
}

static unsigned int
argmax (float*       v,
        unsigned int n)
{
  unsigned int rv = 0;
  for (unsigned int i = 1; i < n; ++i)
    {
      if (v[i] > v[rv])
        {
          rv = i;
        }
    }

  return rv;
}

static char*
to_string (unsigned int* x,
           unsigned int  nx)
{
  static char buf[10240];
  off_t off = 0;

  buf[0] = '\0';

  for (unsigned int n = 0; n < nx; ++n)
    {
      off += sprintf (buf + off, "%s%u", ((n == 0) ? "" : ","), x[n]);
    }

  return buf;
}

static const float* global_pz;

static int
pz_compar_descending (const void* a,
                      const void* b)
{
  const unsigned int* ia = a;
  const unsigned int* ib = b;

  if (global_pz[*ia] < global_pz[*ib])
    {
      return 1;
    }
  else if (global_pz[*ia] > global_pz[*ib])
    {
      return -1;
    }
  else
    {
      return *ia - *ib;
    }
}

static void
print_pz (FILE*            fp,
          const float*     pz,
          MultiModel*      model)
{
  unsigned int index[model->n_labels];

  for (unsigned int j = 0; j < model->n_labels; ++j)
    {
      index[j] = j;
    }

  global_pz = pz;

  qsort (index, model->n_labels, sizeof (*index), pz_compar_descending);

  for (unsigned int j = 0; j < model->n_labels; ++j)
    {
      unsigned int labels[model->n_raw_labels];
      unsigned int depth = index_to_set (index[j], labels, model->n_raw_labels);

      fprintf (fp, "\t%s:%f", to_string (labels, depth), pz[index[j]]);
    }
}

static void
read_loop (FILE*        in,
           MultiModel*  model,
           uint64_t     t,
           float        eta_unscaled,
           float        rho,
           FILE*        predict,
           bool         quiet,
           bool         test_only)
{
  static const int max_ratings = 1024;
  Rating ratings[max_ratings];

  static const int max_tag = 1024;
  char tag[max_tag];

  double since_last_q = 0;
  double total_q = 0;
  unsigned long long bad_lines_since_last = 0;
  unsigned long long empty_lines_since_last = 0;
  unsigned long long example_count = 0;
  unsigned long long rating_count = 0;
  unsigned long long last_bad_lineno = 0;
  unsigned long long last_empty_lineno = 0;
  unsigned long long output_number = 0;
  unsigned long long since_last_example_count = 0;
  unsigned long long since_last_rating_count = 0;

  uint64_t* last_t = (uint64_t*) malloc (model->n_workers * sizeof (uint64_t));
  for (unsigned int n = 0; n < model->n_workers; ++n)
    {
      last_t[n] = t;
    }
  
  float* dloga = (float*) malloc (model->n_workers * model->rank * sizeof (float));
  float dc[model->n_labels * model->rank];
  float dd[model->n_labels * model->rank];

  float eta = eta_unscaled * fastpow (t, rho);

  float pz[model->n_labels];
  int clamp = -1;
  int n_ratings = 0;

  maybe_fprintf (stderr, "cumul     since       example   current   current   current\n"
                         "avg q     last        counter     label   predict   ratings\n");

  for (unsigned long long lineno = 0; ! feof (in); ++lineno)
    {
      float logbeta;

      n_ratings = parse_multi_ratings (in,
                                       ratings,
                                       max_ratings,
                                       tag,
                                       max_tag,
                                       &clamp,
                                       lineno,
                                       model->n_raw_labels,
                                       model->max_raw_labels,
                                       model->n_workers);

      if (n_ratings > 0)
        {
          float q = do_block (model,
                              t,
                              eta,
                              rho,
                              last_t,
                              dloga,
                              dc,
                              dd,
                              &logbeta,
                              pz,
                              clamp,
                              ratings,
                              n_ratings,
                              test_only);

          if (! test_only) { ++t; }

          if (predict)
            {
              fprintf (predict, "%s\t%f", tag, logbeta);
              print_pz (predict, pz, model);
              fprintf (predict, "\n");
            }

          since_last_rating_count += n_ratings;
          since_last_q += q;
          ++since_last_example_count;
        }
      else if (n_ratings == 0)
        {
          /* NB: if output, then output prior */
          ++empty_lines_since_last;
          last_empty_lineno = lineno;

          if (predict)
            {
              /* NB: assuming logbeta prior mean is 0 */
              fprintf (predict, "%s\t0", tag);

              for (unsigned int i = 0; i < model->n_labels; ++i)
                {
                  fprintf (predict, "\t%f", fastexp (model->logpriorz[i]));
                }
              fprintf (predict, "\n");
            }
        }
      else if (n_ratings < 0 && ! feof (in))
        {
          ++bad_lines_since_last;
          last_bad_lineno = lineno;
        }

      if (since_last_example_count > (1ULL << output_number))
        {
          example_count += since_last_example_count;
          rating_count += since_last_rating_count;
          total_q += since_last_q;

          maybe_fprintf (stderr, 
                         "%-9.6f %-9.6f %9llu %9d %9u %9d\n",
                         total_q / rating_count,
                         since_last_q / since_last_rating_count,
                         example_count,
                         clamp,
                         argmax (pz, model->n_labels),
                         n_ratings);

          if (bad_lines_since_last)
            {
              maybe_fprintf (stderr,
                             "warning: encountered %llu malformed input lines, most recent at line %llu\n",
                             bad_lines_since_last,
                             last_bad_lineno);

              bad_lines_since_last = 0;
              last_bad_lineno = 0;
            }

          if (empty_lines_since_last)
            {
              maybe_fprintf (stderr, 
                             "warning: encountered %llu items without ratings, "
                             "most recent at line %llu\n",
                             empty_lines_since_last,
                             last_empty_lineno);

              empty_lines_since_last = 0;
              last_empty_lineno = 0;
            }

          since_last_example_count = 0;
          since_last_rating_count = 0;
          since_last_q = 0;
          ++output_number;
        }
    }

  if (since_last_example_count > 0)
    {
      example_count += since_last_example_count;
      rating_count += since_last_rating_count;
      total_q += since_last_q;

      maybe_fprintf (stderr, 
                     "%-9.6f %-9.6f %9llu %9d %9u %9d\n",
                     total_q / rating_count,
                     since_last_q / since_last_rating_count,
                     example_count,
                     clamp,
                     argmax (pz, model->n_labels),
                     n_ratings);

      if (bad_lines_since_last)
        {
          fprintf (stderr,
                   "error: encountered %llu malformed input lines, "
                   "most recent at line %llu\n",
                   bad_lines_since_last,
                   last_bad_lineno);

          bad_lines_since_last = 0;
          last_bad_lineno = 0;
        }

      if (empty_lines_since_last)
        {
          maybe_fprintf (stderr,
                         "warning: encountered %llu items without ratings, "
                         "most recent at line %llu\n",
                         empty_lines_since_last,
                         last_empty_lineno);

          empty_lines_since_last = 0;
          last_empty_lineno = 0;
        }

      since_last_example_count = 0;
      since_last_rating_count = 0;
      since_last_q = 0;
    }

  if (! test_only)
    {
      maybe_fprintf (stderr, "applying deferred prior updates ...");
      finalize (t, eta, rho, last_t,
                model->n_workers, model->n_items, model->loga, 
                model->gamma, model->rank, model->priorloga);
      maybe_fprintf (stderr, " finished\n");
    }

  free (dloga);
  free (last_t);
}

static void
usage (void)
{
  fprintf (stderr, 
           "usage: multionlineextract [--] [--initial_t t] [--eta eta] [--rho rho] [--n_items items] [--n_raw_labels n_raw_labels] [--max_raw_labels max_raw_labels] [--n_worker_bits worker_bits] [--priorz prior] [--rank rank] [-h] [-q] [-t] [--model modelfile] [--data datafile] [--predict predictfile] [--seed seed]\n"
           "\n"
           "  --: interpret subsequent -whatever as a file argument\n"
           "\n"
           "  -h: display this message.\n"
           "\n"
           "  -q: quiet (don't display progress).\n"
           "\n"
           "  -t: test_only.  do not train.\n"
           "\n"
           "  --initial_t: specify initial time for learning rate decay. (default: 1000)\n"
           "\n"
           "  --eta: specify initial learning rate. (default: 1.0)\n"
           "\n"
           "  --rho: specify learning rate decay exponent. (default: 0.5)\n"
           "\n"
           "  --n_items: specify effective number of items.\n"
           "\n"
           "  --n_raw_labels: specify number of (original) labels.\n"
           "      power sets (subsets) will be formed from these labels\n"
           "\n"
           "  --max_raw_labels: specify maximum size of subset.\n"
           "      this leads to an effective number of labels in the induced subproblem of\n"
           "      Sum[Binomial[n_raw_labels, k], { k, 0, max_raw_labels }]\n"
           "\n"
           "  --n_worker_bits: specify log_2 effective number of workers.  (default: 16)\n"
           "\n"
           "  --rank: specify approximation rank (required).\n"
           "\n"
           "  --priorz: specify (unnormalized) prior distribution as comma-separated float array.\n"
           "\n"
           "  --model: specify model file name.\n"
           "\n"
           "  --data: specify file to read input data from.\n"
           "      if datafile eq '-' use standard input.\n"
           "      if not specified, use standard input.\n"
           "\n"
           "  --predict: specify file to write predictions to.\n"
           "      if not specified, do not write predictions.\n"
           "\n"
           "  --seed: set random seed (default: 45).\n"
           "\n"
          );

  exit (1);
}

static uint64_t
parse_uint64 (char*    s,
              uint64_t min)
{
  char* endptr;
  unsigned long long val = strtoull (s, &endptr, 0);

  if (*s == '\0' || *endptr != '\0')
    {
      fprintf (stderr,
               "error: invalid 64 bit integral constant '%s'\n",
               s);

      exit (1);
    }

  if (val < min)
    {
      fprintf (stderr,
               "error: integer %llu is less than minimum allowed value %llu.\n",
               (long long unsigned int) val,
               (long long unsigned int) min);
      
      exit (1);
    }

  return val;
}

static float
parse_float (char* s,
             float min,
             float max)
{
  char* endptr;
  float val = strtof (s, &endptr);

  if (*s == '\0' || *endptr != '\0')
    {
      fprintf (stderr,
               "error: invalid floating point constant '%s'\n",
               s);

      exit (1);
    }

  if (val < min || val > max)
    {
      fprintf (stderr,
               "error: floating point value %f not within allowed range [%f, %f]\n",
               val,
               min,
               max);

      exit (1);
    }

  return val;
}

static unsigned int
parse_unsigned_int_arg (char*        s,
                        unsigned int min,
                        unsigned int max)
{
  char* endptr;
  unsigned int val = strtoul (s, &endptr, 0);

  if (*s == '\0' || *endptr != '\0')
    {
      fprintf (stderr,
               "error: invalid integral constant '%s'\n",
               s);

      exit (1);
    }

  if (val < min || val > max)
    {
      fprintf (stderr,
               "error: integral value %u not within allowed range [%u, %u]\n",
               val,
               min,
               max);

      exit (1);
    }

  return val;
}

static void
read_logpriorz (char*        s,
                float*       logpriorz,
                unsigned int n_labels)
{
  float sum = 0;

  for (unsigned int n = 0; n < n_labels; ++n)
    {
      char* comma = strchr (s, ',');

      if (comma == NULL && n + 1 < n_labels)
        {
          fprintf (stderr, "not enough parameters for logpriorz\n");
          exit (1);
        }

      if (comma) { *comma = '\0'; }

      logpriorz[n] = parse_float (s, 0, 1U << 31);
      sum += logpriorz[n];

      s = comma + 1;
    }

  for (unsigned int n = 0; n < n_labels; ++n)
    {
      logpriorz[n] /= sum;
      logpriorz[n] = fastlog (logpriorz[n]);
    }
}

static void
print_logpriorz (const float* logpriorz,
                 unsigned int n_labels)
{
  fprintf (stderr, "priorz = ");
  for (unsigned int n = 0; n < n_labels; ++n)
    {
      fprintf (stderr, "%s%f", (n == 0) ? "" : ",", fastexp (logpriorz[n]));
    }

  fprintf (stderr, "\n");
}

static void
print_gamma (const float* gamma,
             unsigned int rank)
{
  fprintf (stderr, "\ngamma:");

  for (unsigned int n = 0; n < rank; ++n)
    {
      fprintf (stderr, "\t%4.4f", gamma[n]);
    }

  fprintf (stderr, "\n");
}

int 
main (int       argc, 
      char*     argv[])
{
  const char* dataname = NULL;
  const char* modelname = NULL;
  const char* predictname = NULL;
  uint64_t initial_t = 1000;
  float eta = 1.0;
  float rho = 0.5;
  unsigned int n_items = 0;
  unsigned int n_raw_labels = 0;
  unsigned int max_raw_labels = 0;
  unsigned int n_workers = 1U << 16;
  unsigned int rank = 0;
  unsigned int seed = 0;
  float* logpriorz = NULL;

  bool quiet = false;
  bool test_only = false;

  (void) parse_ratings;

  if (argc == 1)
    {
      usage ();
    }

  while (argc > 1)
    {
      if (strcmp (argv[1], "--") == 0)
        {
          ++argv;
          --argc;
          break;
        }
      else if (strcmp (argv[1], "-h") == 0)
        {
          usage ();
        }
      else if (strcmp (argv[1], "--initial_t") == 0 && argc > 2)
        {
          ++argv;
          --argc;

          initial_t = parse_uint64 (argv[1], 1);
          ++argv;
          --argc;
        }
      else if (strcmp (argv[1], "--eta") == 0 && argc > 2)
        {
          ++argv;
          --argc;

          eta = parse_float (argv[1], 0, 10);
          ++argv;
          --argc;
        }
      else if (strcmp (argv[1], "--rho") == 0 && argc > 2)
        {
          ++argv;
          --argc;

          rho = parse_float (argv[1], 0, 0.99);
          ++argv;
          --argc;
        }
      else if (strcmp (argv[1], "--n_items") == 0 && argc > 2)
        {
          ++argv;
          --argc;

          n_items = parse_unsigned_int_arg (argv[1], 1, 1U << 31);
          ++argv;
          --argc;
        }
      else if (strcmp (argv[1], "--n_raw_labels") == 0 && argc > 2)
        {
          ++argv;
          --argc;

          n_raw_labels = parse_unsigned_int_arg (argv[1], 1, 64);
          ++argv;
          --argc;
        }
      else if (strcmp (argv[1], "--max_raw_labels") == 0 && argc > 2)
        {
          ++argv;
          --argc;

          if (n_raw_labels == 0)
            {
              fprintf (stderr,
                       "error: have to specify n_raw_labels before specifying max_raw_labels\n");
              exit (1);
            }

          max_raw_labels = parse_unsigned_int_arg (argv[1], 1, n_raw_labels);
          ++argv;
          --argc;
        }
      else if (strcmp (argv[1], "--n_worker_bits") == 0 && argc > 2)
        {
          ++argv;
          --argc;

          unsigned int n_worker_bits = parse_unsigned_int_arg (argv[1], 8, 32);
          n_workers = 1U << n_worker_bits;
          ++argv;
          --argc;
        }
      else if (strcmp (argv[1], "--priorz") == 0 && argc > 2)
        {
          ++argv;
          --argc;

          if (n_raw_labels == 0)
            {
              fprintf (stderr,
                       "error: have to specify n_raw_labels before specifying prior distribution\n");
              exit (1);
            }

          logpriorz = alloca (n_raw_labels * sizeof (float));
          read_logpriorz (argv[1], logpriorz, n_raw_labels);
          ++argv;
          --argc;
        }
      else if (! rank && strcmp (argv[1], "--rank") == 0)
        {
          ++argv;
          --argc;

          rank = parse_unsigned_int_arg (argv[1], 1, 1024);
          ++argv;
          --argc;
        }
      else if (! quiet && strcmp (argv[1], "-q") == 0)
        {
          quiet = true;
          ++argv;
          --argc;
        }
      else if (! test_only && strcmp (argv[1], "-t") == 0)
        {
          test_only = true;
          ++argv;
          --argc;
        }
      else if (strcmp (argv[1], "--predict") == 0 && argc > 2)
        {
          ++argv;
          --argc;

          predictname = argv[1];
          ++argv;
          --argc;
        }
      else if (strcmp (argv[1], "--model") == 0 && argc > 2)
        {
          ++argv;
          --argc;

          modelname = argv[1];
          ++argv;
          --argc;
        }
      else if (strcmp (argv[1], "--data") == 0 && argc > 2)
        {
          ++argv;
          --argc;

          dataname = argv[1];
          ++argv;
          --argc;
        }
      else if (strcmp (argv[1], "--seed") == 0 && argc > 2)
        {
          ++argv;
          --argc;

          seed = parse_unsigned_int_arg (argv[1], 0, 1U << 31);
          ++argv;
          --argc;
        }
      else
        {
          fprintf (stderr, "error: unknown argument '%s'\n", argv[1]);
          exit (1);
        }
    }

  if (argc > 1)
    {
      usage ();
    }

  if (modelname == NULL)
    {
      fprintf (stderr, "error: model file not specified\n");
      exit (1);
    }

  if (seed == 0)
    {
      seed = 45;
    }

  srand48 (seed);

  MultiModel* model;
  struct stat buf;

  if (stat (modelname, &buf) == 0)
    {
      model = multi_model_open (modelname, test_only);
    }
  else
    {
      if (n_items == 0)
        {
          fprintf (stderr, "error: n_items not specified\n");
          exit (1);
        }

      if (n_raw_labels == 0)
        {
          fprintf (stderr, "error: n_raw_labels not specified\n");
          exit (1);
        }

      if (max_raw_labels == 0)
        {
          fprintf (stderr, "error: max_raw_labels not specified\n");
          exit (1);
        }

      if (rank == 0)
        {
          fprintf (stderr, "error: rank not specified\n");
          exit (1);
        }

      if (logpriorz == NULL)
        {
          fprintf (stderr, "error: priorz not specified\n");
          exit (1);
        }

      model = multi_model_creat (modelname,
                                       n_items,
                                       n_raw_labels,
                                       max_raw_labels,
                                       n_workers,
                                       rank,
                                       logpriorz);
    }

  if (model == NULL)
    {
      fprintf (stderr, "error: failed to open model file\n");
      exit (1);
    }

  FILE* dfp = 
    (dataname == NULL || strcmp (dataname, "-") == 0) 
      ? stdin 
      : fopen (dataname, "r");

  if (dfp == NULL)
    {
      fprintf (stderr, 
               "fatal: can't open %s: %s\n",
               (dataname == NULL) ? "stdin" : dataname,
               strerror (errno));
      exit (1);
    }

  FILE* predict =
    (predictname == NULL) ? NULL :
    (strcmp (predictname, "-") == 0) ? stdout : fopen (predictname, "w");

  if (predict == NULL && predictname != NULL)
    {
      fprintf (stderr, 
               "fatal: can't open %s: %s\n",
               predictname,
               strerror (errno));
      exit (1);
    }

  char etabuf[1024];
  char rhobuf[1024];

  maybe_fprintf (stderr,
                 "seed = %u\n"
                 "initial_t = %llu\n"
                 "eta = %s\n"
                 "rho = %s\n"
                 "n_items = %u\n"
                 "n_raw_labels = %u\n"
                 "max_raw_labels = %u\n"
                 "n_labels (induced) = %u\n"
                 "n_workers = %u\n"
                 "rank = %u\n"     
                 "test_only = %s\n"
                 "prediction file = %s\n",
                 seed,
                 (long long unsigned int) initial_t,
                 test_only ? "(n/a)" : (snprintf (etabuf,
                                                  sizeof (etabuf),
                                                  "%f",
                                                  eta),
                                        etabuf),
                 test_only ? "(n/a)" : (snprintf (rhobuf,
                                                  sizeof (rhobuf),
                                                  "%f",
                                                  rho),
                                        rhobuf),
                 model->n_items,
                 model->n_raw_labels,
                 model->max_raw_labels,
                 model->n_labels,
                 model->n_workers,
                 model->rank,
                 test_only ? "true" : "false",
                 (predictname == NULL) ? "(no output)" :
                   (strcmp (predictname, "-") == 0) ? "stdout" : predictname);

  if (! quiet)
    {
      print_logpriorz (model->rawlogpriorz, model->n_raw_labels);
    }

  read_loop (dfp, model, initial_t, eta, rho, predict, quiet, test_only);

  if (! quiet && ! test_only)
    {
      print_gamma (model->gamma, model->rank);
    }

  multi_model_close (model);

  return 0;
}
